# {0: 'Agaricus', 1: 'Amanita', 2: 'Boletus', 3: 'Cortinarius', 4: 'Entoloma', 5: 'Hygrocybe', 6: 'Lactarius', 7: 'Russula', 8: 'Suillus'}
import time

import cv2
import numpy as np
import os
from sklearn.svm import SVC
from skimage.feature import hog
import joblib


# 图像预处理函数：调整图像大小，灰度化，归一化
def preprocess_image(image, target_size=(128, 128)):
    image_resized = cv2.resize(image, target_size)
    gray_image = cv2.cvtColor(image_resized, cv2.COLOR_BGR2GRAY)
    normalized_image = gray_image / 255.0
    return normalized_image


# 提取图像的 HOG 特征
def extract_hog_features(images):
    hog_features = []
    for image in images:
        hog_feature = hog(image, orientations=9, pixels_per_cell=(8, 8), cells_per_block=(2, 2), transform_sqrt=True)
        hog_features.append(hog_feature)
    return np.array(hog_features)


# 加载模型
def load_model(model_path):
    return joblib.load(model_path)


# 保存模型
def save_model(model, model_path):
    joblib.dump(model, model_path)


# 训练模型
def train_model(X_train, y_train):
    svm_model = SVC(kernel='linear')
    svm_model.fit(X_train, y_train)
    print("SVM model trained.")
    return svm_model


# 加载已训练好的模型和标签映射
def load_model_and_labels(model_path, label_map):
    svm_model = load_model(model_path)
    return svm_model, label_map


# 植物图像识别
def predict_plant(image_path, svm_model, label_map):
    image = cv2.imread(image_path)
    preprocessed_image = preprocess_image(image)
    hog_feature = extract_hog_features([preprocessed_image])
    prediction = svm_model.predict(hog_feature)
    plant_label = label_map[prediction[0]]
    return plant_label


# 加载训练数据
def load_data(folder_path):
    images = []
    labels = []
    label_map = {}
    for label, plant_folder in enumerate(os.listdir(folder_path)):
        label_map[label] = plant_folder
        for image_name in os.listdir(os.path.join(folder_path, plant_folder)):
            image_path = os.path.join(folder_path, plant_folder, image_name)
            image = cv2.imread(image_path)
            preprocessed_image = preprocess_image(image)
            images.append(preprocessed_image)
            labels.append(label)
    return images, labels, label_map

# Agaricus - 伞菇属
# Amanita - 鹅膏菌属
# Boletus - 牛肝菌属
# Cortinarius - 腰带菌属
# Entoloma - 桔黄褶菌属
# Hygrocybe - 露脐菌属
# Lactarius - 乳菇属
# Russula - 红菇属
# Suillus - 双孢菌属
if __name__ == "__main__":
    t0 = time.time()
    # 加载已有模型和标签映射
    model_path = "model/svm_model.joblib" # 模型路径
    label_map = {0: 'Agaricus', 1: 'Amanita', 2: 'Boletus', 3: 'Cortinarius', 4: 'Entoloma', 5: 'Hygrocybe',
                 6: 'Lactarius', 7: 'Russula', 8: 'Suillus'}    # 标签映射
    svm_model, label_map = load_model_and_labels(model_path, label_map)

    # 测试植物图像识别
    test_image_path = "img.png"  # 测试图像路径
    predicted_plant = predict_plant(test_image_path, svm_model, label_map)
    print(f"Predicted plant: {predicted_plant}")
    t1 = time.time()
    elapsed_time = t1 - t0
    print(f'耗时是: {elapsed_time:.6f} 秒')